package com.sqs.limiter;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.util.concurrent.RateLimiter;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

@Component
public class IpRateLimiter {

    // 缓存每个 key 的 RateLimiter，带自动过期机制
    private final Cache<String, RateLimiter> limiterCache = CacheBuilder.newBuilder()
            .expireAfterAccess(10, TimeUnit.MINUTES) // 10分钟不访问就清除
            .maximumSize(1000) // 防止缓存过大
            .build();

    /**
     * 尝试获取令牌（非阻塞）
     * @param key 唯一限流 key（如 IP、用户ID、接口名）
     * @param permitsPerSecond 限流速率（每秒允许多少请求）
     * @return 是否允许通过
     */
    public boolean tryAcquire(String key, double permitsPerSecond) {
        RateLimiter rateLimiter = limiterCache.getIfPresent(key);
        if (rateLimiter == null) {
            synchronized (this) {
                rateLimiter = limiterCache.getIfPresent(key);
                if (rateLimiter == null) {
                    rateLimiter = RateLimiter.create(permitsPerSecond);
                    limiterCache.put(key, rateLimiter);
                }
            }
        }
        return rateLimiter.tryAcquire(); // 非阻塞获取令牌
    }

}
